#ifndef TIEC_TREE_TRANSFORMER_H_
#define TIEC_TREE_TRANSFORMER_H_

#include "tree.h"

namespace NS_TIEC {
  template <typename Param>
  class TreeTransformer : public TreeVisitor<PTR<Tree>, Param> {
  public:
    template <typename T, typename = std::enable_if<std::is_base_of_v<Tree, T>>>
    PTR<T> transform(const PTR<T>& tree, Param param) {
      const PTR<Tree> r = TreeVisitor<PTR<Tree>, Param>::visitTree(tree, param);
      return std::static_pointer_cast<T>(r);
    }

    template <typename T, typename = std::enable_if<std::is_base_of_v<Tree, T>>>
    List<PTR<T>> transform(List<PTR<T>>& trees, Param param) {
      for (PTR<T>& tree : trees) {
        tree = transform(tree, param);
      }
      return trees;
    }

    PTR<Tree> visitCompilationUnitTree(const PTR<CompilationUnitTree>& tree, Param param) override {
      tree->defs = transform(tree->defs, param);
      return tree;
    }

    PTR<Tree> visitClassTree(const PTR<ClassTree>& tree, Param param) override {
      tree->members = transform(tree->members, param);
      tree->codes = transform(tree->codes, param);
      return tree;
    }

    PTR<Tree> visitVariableTree(const PTR<VariableTree>& tree, Param param) override {
      tree->initializer = transform(tree->initializer, param);
      return tree;
    }

    PTR<Tree> visitMethodTree(const PTR<MethodTree>& tree, Param param) override {
      tree->parameters = transform(tree->parameters, param);
      tree->body = transform(tree->body, param);
      return tree;
    }

    PTR<Tree> visitEventTree(const PTR<EventTree>& tree, Param param) override {
      tree->event_obj = transform(tree->event_obj, param);
      tree->parameters = transform(tree->parameters, param);
      tree->body = transform(tree->body, param);
      return tree;
    }

    PTR<Tree> visitAnnotationTree(const PTR<AnnotationTree>& tree, Param param) override {
      tree->args = transform(tree->args, param);
      return tree;
    }

    PTR<Tree> visitCodeStatementTree(const PTR<CodeStatementTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitCodeExpressionTree(const PTR<CodeExpressionTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitCodeSegmentTree(const PTR<CodeSegmentTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitBlockTree(const PTR<BlockTree>& tree, Param param) override {
      tree->statements = transform(tree->statements, param);
      return tree;
    }

    PTR<Tree> visitExpressionStatementTree(const PTR<ExpressionStatementTree>& tree, Param param) override {
      tree->expr = transform(tree->expr, param);
      return tree;
    }

    PTR<Tree> visitIfTree(const PTR<IfTree>& tree, Param param) override {
      tree->condition = transform(tree->condition, param);
      tree->then_part = transform(tree->then_part, param);
      tree->else_part = transform(tree->else_part, param);
      return tree;
    }

    PTR<Tree> visitSwitchTree(const PTR<SwitchTree>& tree, Param param) override {
      tree->selector = transform(tree->selector, param);
      tree->cases = transform(tree->cases, param);
      tree->default_case = transform(tree->default_case, param);
      return tree;
    }
    
    PTR<Tree> visitCaseTree(const PTR<CaseTree>& tree, Param param) override {
      tree->conditions = transform(tree->conditions, param);
      tree->body = transform(tree->body, param);
      return tree;
    }

    PTR<Tree> visitForLoopTree(const PTR<ForLoopTree>& tree, Param param) override {
      tree->target = transform(tree->target, param);
      tree->start = transform(tree->start, param);
      tree->end = transform(tree->end, param);
      tree->step = transform(tree->step, param);
      tree->body = transform(tree->body, param);
      return tree;
    }

    PTR<Tree> visitForEachTree(const PTR<ForEachTree>& tree, Param param) override {
      tree->source = transform(tree->source, param);
      tree->body = transform(tree->body, param);
      return tree;
    }

    PTR<Tree> visitWhileLoopTree(const PTR<WhileLoopTree>& tree, Param param) override {
      tree->condition = transform(tree->condition, param);
      tree->body = transform(tree->body, param);
      return tree;
    }

    PTR<Tree> visitLiteralTree(const PTR<LiteralTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitArrayTypeTree(const PTR<ArrayTypeTree>& tree, Param param) override {
      tree->expr = transform(tree->expr, param);
      tree->size = transform(tree->size, param);
      return tree;
    }

    PTR<Tree> visitTemplateTypeTree(const PTR<TemplateTypeTree>& tree, Param param) override {
      tree->class_name = transform(tree->class_name, param);
      tree->args = transform(tree->args, param);
      return tree;
    }

    PTR<Tree> visitIdentifierTree(const PTR<IdentifierTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitFieldAccessTree(const PTR<FieldAccessTree>& tree, Param param) override {
      tree->prefix = transform(tree->prefix, param);
      return tree;
    }

    PTR<Tree> visitMethodInvokeTree(const PTR<MethodInvokeTree>& tree, Param param) override {
      tree->method = transform(tree->method, param);
      tree->args = transform(tree->args, param);
      return tree;
    }

    PTR<Tree> visitNewClassTree(const PTR<NewClassTree>& tree, Param param) override {
      tree->args = transform(tree->args, param);
      return tree;
    }

    PTR<Tree> visitNewArrayTree(const PTR<NewArrayTree>& tree, Param param) override {
      tree->elements = transform(tree->elements, param);
      return tree;
    }

    PTR<Tree> visitArrayAccessTree(const PTR<ArrayAccessTree>& tree, Param param) override {
      tree->selected = transform(tree->selected, param);
      tree->index = transform(tree->index, param);
      return tree;
    }

    PTR<Tree> visitBreakTree(const PTR<BreakTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitContinueTree(const PTR<ContinueTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitReturnTree(const PTR<ReturnTree>& tree, Param param) override {
      tree->value = transform(tree->value, param);
      return tree;
    }

    PTR<Tree> visitSubscribeEventTree(const PTR<SubscribeEventTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitTypeCastTree(const PTR<TypeCastTree>& tree, Param param) override {
      tree->expr = transform(tree->expr, param);
      tree->type_expr = transform(tree->type_expr, param);
      return tree;
    }

    PTR<Tree> visitAssignTree(const PTR<AssignTree>& tree, Param param) override {
      tree->left = transform(tree->left, param);
      tree->right = transform(tree->right, param);
      return tree;
    }

    PTR<Tree> visitExchangeTree(const PTR<ExchangeTree>& tree, Param param) override {
      tree->left = transform(tree->left, param);
      tree->right = transform(tree->right, param);
      return tree;
    }

    PTR<Tree> visitInstanceOfTree(const PTR<InstanceOfTree>& tree, Param param) override {
      tree->expr = transform(tree->expr, param);
      tree->type_expr = transform(tree->type_expr, param);
      return tree;
    }

    PTR<Tree> visitMultiExprTree(const PTR<MultiExprTree>& tree, Param param) override {
      tree->expressions = transform(tree->expressions, param);
      return tree;
    }

    PTR<Tree> visitAwaitTree(const PTR<AwaitTree>& tree, Param param) override {
      tree->expr = transform(tree->expr, param);
      return tree;
    }

    PTR<Tree> visitParensTree(const PTR<ParensTree>& tree, Param param) override {
      tree->expr = transform(tree->expr, param);
      return tree;
    }

    PTR<Tree> visitUnaryTree(const PTR<UnaryTree>& tree, Param param) override {
      tree->expr = transform(tree->expr, param);
      return tree;
    }

    PTR<Tree> visitBinaryTree(const PTR<BinaryTree>& tree, Param param) override {
      tree->left = transform(tree->left, param);
      tree->right = transform(tree->right, param);
      return tree;
    }

    PTR<Tree> visitCodeTextTree(const PTR<CodeTextTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitCodeRefTree(const PTR<CodeRefTree>& tree, Param param) override {
      return tree;
    }

    PTR<Tree> visitErroneousTree(const PTR<ErroneousTree>& tree, Param param) override {
      return tree;
    }
  };
}

#endif //TIEC_TREE_TRANSFORMER_H_
